import onnx
import yaml, sys
import numpy as np
from yaml.loader import SafeLoader
from . import utils

def parse_yaml(yaml_file):
    alpha_list = []
    control_list = []

    # Open the file and load the file
    with open(yaml_file) as f:
        data = yaml.load(f, Loader=SafeLoader)
        if 'postproc' in data.keys():
            print('got postproc---')
            postproc_list = data['postproc']
            for p in postproc_list:
                if 'alpha' in p.keys():
                    alpha_list_ = p['alpha']
                    #if len(alpha_list_) == 3:
                    alpha_list = alpha_list_
                    print('got alpha values:', alpha_list) 
                
                if 'control' in p.keys():
                    control_list = p['control']
                    print('got control values:', control_list)

    if len(control_list) == 0:
        return {}

    postproc_dict = {} 

    if len(alpha_list) > 0:
        postproc_dict['alpha'] = alpha_list

    postproc_dict['control'] = control_list

    return postproc_dict    

def get_const_type(val):
    types = {
        'const_alpha' : 2,
        'const_control_post' : 6
    }

    return types.get(val, None)

def get_init_val(init):
    val = []

    if init.data_type != get_const_type(init.name):
        print('Error, mismatch init type', get_const_type(init.name), init.data_type, init.name)
        return False

    dtype = init.data_type
    np_dtype = utils.convert_ort_type_2_np(dtype)
    if init.raw_data:
        params_list = np.fromstring(init.raw_data, dtype=np_dtype)
        for p in params_list:
            val.append(p)
    else:
        data_list = utils.get_data_list(dtype, init)
        for p in data_list:
            val.append(p)

    return val        

def get_snr(gpu_array, cpu_array):
    diff_array = np.subtract(cpu_array, gpu_array)
    x = np.square(diff_array)
    x = np.sum(x)

    y = np.square(cpu_array)
    y = np.sum(y) 

    snr = (x) / (y + 1e-7)

    snr = np.mean(snr)

    #print('snr:', snr)

    return snr  

def analyze(onnxfile, **kwargs):
    yaml = kwargs['yaml']

    preproc_dict = parse_yaml(yaml) 

    model = onnx.load(onnxfile)

    expect_input = ['const_alpha', 'const_control_post']

    first_check = False
    for node in model.graph.node:
        if node.op_type == 'PostProc':
            print('Got PostProc Node, input:', node.input)
            if set(expect_input).issubset(set(node.input)):
                #print('fist check is ok')
                first_check = True
                break
            else:
                return False

    if first_check == True:
        init_list = []
        for init in model.graph.initializer:
            init_list.append(init.name)

        if set(expect_input).issubset(set(init_list)) == False:
            print('Error: expect_input is not in init_list')
            return False

        d = {}
        for init in model.graph.initializer:
            if init.name in expect_input:
                val = get_init_val(init)
                #print('get val:', init.name, val)
                d[init.name] = val

        for k, v in preproc_dict.items():
                key = 'const_' + k
                if k == 'control':
                    key = key + '_post'

                if key in d.keys():
                    print('kv', k, ':', v)
                    np_array_1 = np.array(v).flatten()
                    np_array_2 = np.array(d[key]).flatten()
                    snr = get_snr(np_array_1, np_array_2)
                    if snr > 0.1:
                        print('ERROR:snr too large')
                        return False

        return True    

    return False